#!/usr/bin/env python3
import rclpy
from rclpy.node import Node
from sensor_msgs.msg import PointCloud2
import sensor_msgs_py.point_cloud2 as pc2
import numpy as np
import open3d as o3d
import threading
from std_srvs.srv import Trigger
import os
from datetime import datetime

class PointCloudSaver(Node):
    def __init__(self):
        super().__init__('pointcloud_saver')
        
        # 订阅RTAB-Map的点云话题
        self.subscription = self.create_subscription(
            PointCloud2,
            '/rtabmap/cloud_map',
            self.pointcloud_callback,
            10)
        
        self.pointcloud_data = None
        self.lock = threading.Lock()
        
        self.get_logger().info('点云保存节点已启动，等待数据...')
        
        # 创建保存服务
        self.save_service = self.create_service(
            Trigger,
            'save_pointcloud',
            self.save_callback)
        
        # 创建自动保存目录
        self.save_dir = os.path.expanduser('~/pointcloud_maps')
        os.makedirs(self.save_dir, exist_ok=True)
    
    def pointcloud_callback(self, msg):
        with self.lock:
            try:
                # 转换点云数据
                points = pc2.read_points(msg, field_names=("x", "y", "z", "rgb"), skip_nans=True)
                points_list = list(points)
                
                if len(points_list) > 0:
                    # 转换为numpy数组
                    self.pointcloud_data = np.array(points_list)
                    self.get_logger().info(f'收到点云数据: {len(self.pointcloud_data)} 个点')
                else:
                    self.get_logger().warn('收到的点云数据为空')
                    
            except Exception as e:
                self.get_logger().error(f'处理点云数据时出错: {str(e)}')
    
    def save_callback(self, request, response):
        with self.lock:
            if self.pointcloud_data is not None and len(self.pointcloud_data) > 0:
                try:
                    # 创建Open3D点云对象
                    pcd = o3d.geometry.PointCloud()
                    
                    # 设置点坐标
                    points = self.pointcloud_data[:, :3]
                    pcd.points = o3d.utility.Vector3dVector(points)
                    
                    # 设置颜色（如果有RGB信息）
                    if self.pointcloud_data.shape[1] >= 4:
                        colors = self.pointcloud_data[:, 3]
                        # 将RGB整数转换为0-1范围的浮点数
                        rgb_colors = np.zeros((len(colors), 3))
                        for i, rgb in enumerate(colors):
                            r = ((rgb >> 16) & 0xff) / 255.0
                            g = ((rgb >> 8) & 0xff) / 255.0
                            b = (rgb & 0xff) / 255.0
                            rgb_colors[i] = [r, g, b]
                        pcd.colors = o3d.utility.Vector3dVector(rgb_colors)
                    
                    # 生成文件名
                    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
                    output_path = os.path.join(self.save_dir, f'office_map_{timestamp}.ply')
                    
                    # 保存为PLY文件
                    o3d.io.write_point_cloud(output_path, pcd)
                    
                    response.success = True
                    response.message = f'点云已保存到: {output_path}，包含 {len(points)} 个点'
                    self.get_logger().info(response.message)
                    
                except Exception as e:
                    response.success = False
                    response.message = f'保存失败: {str(e)}'
                    self.get_logger().error(response.message)
            else:
                response.success = False
                response.message = '没有可用的点云数据'
                self.get_logger().warn(response.message)
        
        return response

def main():
    rclpy.init()
    node = PointCloudSaver()
    
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        node.get_logger().info('节点被用户中断')
    finally:
        node.destroy_node()
        rclpy.shutdown()

if __name__ == '__main__':
    main()
